{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6b557135",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -q datasets  # tokenizers is optional, we default to a tiny char tokenizer\n",
    "#!tt-smi -r\n",
    "import os, sys, math, random, textwrap\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from dataclasses import dataclass\n",
    "from typing import List, Tuple, Dict, Optional\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Point PYTHONPATH to your TTML build (adjust if needed)\n",
    "sys.path.append(f'{os.environ[\"TT_METAL_HOME\"]}/tt-train/sources/ttml')\n",
    "\n",
    "import ttml\n",
    "\n",
    "# Repro\n",
    "def set_seed(seed: int = 42):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "set_seed(42)\n",
    "\n",
    "@dataclass\n",
    "class TrainConfig:\n",
    "    seq_len: int = 256          # context length\n",
    "    batch_size: int = 64\n",
    "    steps: int = 5000              # total train steps\n",
    "    eval_every: int = 200\n",
    "    lr: float = 3e-4\n",
    "    beta1: float = 0.9\n",
    "    beta2: float = 0.999\n",
    "    eps: float = 1e-8\n",
    "    weight_decay: float = 0.01\n",
    "    n_layer: int = 6\n",
    "    n_head: int = 6\n",
    "    n_embd: int = 384\n",
    "    seed: int = 42\n",
    "    use_bpe: bool = False          # optional: use HF tokenizers BPE instead of char tokenizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ea598840",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 1003853, 111540)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "def load_shakespeare_text() -> str:\n",
    "\n",
    "    ds = load_dataset(\n",
    "        \"text\",\n",
    "        data_files={\"train\": f'{os.environ[\"TT_METAL_HOME\"]}/tt-train/data/shakespeare.txt'},\n",
    "    )\n",
    "    text = \"\\n\".join(ds[\"train\"][\"text\"])\n",
    "    return text\n",
    "# --- Character-level tokenizer (default) ---\n",
    "class CharTokenizer:\n",
    "    def __init__(self, text: str):\n",
    "        # Sorted stable alphabet for reproducibility\n",
    "        chars = sorted(list(set(text)))\n",
    "        self.stoi: Dict[str, int] = {ch: i for i, ch in enumerate(chars)}\n",
    "        self.itos: List[str] = chars\n",
    "\n",
    "    @property\n",
    "    def vocab_size(self) -> int:\n",
    "        return len(self.itos)\n",
    "\n",
    "    def encode(self, s: str) -> List[int]:\n",
    "        return [self.stoi[c] for c in s]\n",
    "\n",
    "    def decode(self, ids: List[int]) -> str:\n",
    "        return \"\".join(self.itos[i] for i in ids)\n",
    "\n",
    "# --- Optional: BPE tokenizer using HF `tokenizers` ---\n",
    "def build_bpe_tokenizer_from_text(tmp_corpus_path: str = \"shakespeare_corpus.txt\"):\n",
    "    # Minimal training pipeline; produces subword tokens (not pure char)\n",
    "    # Only used if cfg.use_bpe=True\n",
    "    from tokenizers import Tokenizer\n",
    "    from tokenizers.models import BPE\n",
    "    from tokenizers.trainers import BpeTrainer\n",
    "    from tokenizers.pre_tokenizers import ByteLevel\n",
    "\n",
    "    tok = Tokenizer(BPE(unk_token=\"[UNK]\"))\n",
    "    tok.pre_tokenizer = ByteLevel()\n",
    "    trainer = BpeTrainer(vocab_size=8192, special_tokens=[\"[UNK]\", \"[BOS]\", \"[EOS]\"])\n",
    "    tok.train(files=[tmp_corpus_path], trainer=trainer)\n",
    "    return tok\n",
    "\n",
    "# --- Build tokenizer + dataset splits ---\n",
    "def prepare_data(cfg: TrainConfig):\n",
    "    text = load_shakespeare_text()\n",
    "\n",
    "    if cfg.use_bpe:\n",
    "        # Train a tiny BPE tokenizer on the fly\n",
    "        with open(\"shakespeare_corpus.txt\", \"w\", encoding=\"utf-8\") as f:\n",
    "            f.write(text)\n",
    "        bpe = build_bpe_tokenizer_from_text(\"shakespeare_corpus.txt\")\n",
    "\n",
    "        def encode_all(t: str) -> np.ndarray:\n",
    "            return np.array(bpe.encode(t).ids, dtype=np.uint32)\n",
    "\n",
    "        def decode(ids: List[int]) -> str:\n",
    "            return bpe.decode(ids)\n",
    "\n",
    "        vocab_size = bpe.get_vocab_size()\n",
    "        encode, decode_fn = encode_all, decode\n",
    "    else:\n",
    "        # Default: fast, deterministic character-level tokenizer\n",
    "        ctok = CharTokenizer(text)\n",
    "        vocab_size = (ctok.vocab_size + 31) // 32 * 32  # pad to multiple of 32\n",
    "        encode = lambda t: np.array(ctok.encode(t), dtype=np.uint32)\n",
    "        decode_fn = lambda ids: ctok.decode(list(ids))\n",
    "\n",
    "    # Encode full corpus, split 90/10\n",
    "    ids = encode(text)\n",
    "    n = len(ids)\n",
    "    n_train = int(n * 0.9)\n",
    "    train_ids = ids[:n_train]\n",
    "    val_ids = ids[n_train:]\n",
    "    return train_ids, val_ids, vocab_size, decode_fn\n",
    "\n",
    "cfg = TrainConfig()\n",
    "train_ids, val_ids, vocab_size, decode = prepare_data(cfg)\n",
    "vocab_size, len(train_ids), len(val_ids)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b68faffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(split_ids: np.ndarray, seq_len: int, batch_size: int) -> Tuple[np.ndarray, np.ndarray]:\n",
    "    n = len(split_ids) - seq_len - 1\n",
    "    ix = np.random.randint(0, n, size=(batch_size,))\n",
    "    x = np.stack([split_ids[i:i+seq_len] for i in ix], axis=0)             # [B, T]\n",
    "    y = np.stack([split_ids[i+1:i+seq_len+1] for i in ix], axis=0)         # [B, T] next-token targets\n",
    "    return x.astype(np.uint32), y.astype(np.uint32)\n",
    "\n",
    "def build_causal_mask(T: int) -> np.ndarray:\n",
    "    # [1,1,T,T] float32 with 1s for allowed positions (i >= j), else 0\n",
    "    m = np.tril(np.ones((T, T), dtype=np.float32))\n",
    "    return m.reshape(1, 1, T, T)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "112034c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model_and_optim(cfg: TrainConfig, vocab_size: int, seq_len: int):\n",
    "    # GPT2 config via your bindings\n",
    "    gcfg = ttml.models.gpt2.GPT2TransformerConfig()\n",
    "    gcfg.num_heads = cfg.n_head\n",
    "    gcfg.embedding_dim = cfg.n_embd\n",
    "    gcfg.num_blocks = cfg.n_layer\n",
    "    gcfg.vocab_size = int(vocab_size)\n",
    "    gcfg.max_sequence_length = seq_len\n",
    "    gcfg.dropout_prob = 0.2\n",
    "    # optional flags exist (runner_type, weight_tying, positional_embedding_type, experimental, ...)\n",
    "    # we keep defaults for a minimal demo\n",
    "\n",
    "    model = ttml.models.gpt2.create_gpt2_model(gcfg)\n",
    "\n",
    "    # AdamW via your bindings\n",
    "    adamw_cfg = ttml.optimizers.AdamWConfig.make(\n",
    "        float(cfg.lr),\n",
    "        float(cfg.beta1),\n",
    "        float(cfg.beta2),\n",
    "        float(cfg.eps),\n",
    "        float(cfg.weight_decay),\n",
    "    )\n",
    "    optim = ttml.optimizers.AdamW(model.parameters(), adamw_cfg)\n",
    "    return model, optim\n",
    "model, optim = create_model_and_optim(cfg, vocab_size, cfg.seq_len)\n",
    "model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b12d32b",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name '_tml' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 52\u001b[0m\n\u001b[1;32m     47\u001b[0m             \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstep \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstep\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m5d\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m | train_loss \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtrain_losses[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m | val_loss \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mval_losses[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m.4f\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m train_losses, val_losses\n\u001b[0;32m---> 52\u001b[0m train_losses, val_losses \u001b[38;5;241m=\u001b[39m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcfg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_ids\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_ids\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;28mlen\u001b[39m(train_losses), \u001b[38;5;28mlen\u001b[39m(val_losses)\n",
      "Cell \u001b[0;32mIn[5], line 3\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(cfg, model, optim, train_ids, val_ids)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtrain\u001b[39m(cfg: TrainConfig, model, optim, train_ids: np\u001b[38;5;241m.\u001b[39mndarray, val_ids: np\u001b[38;5;241m.\u001b[39mndarray):\n\u001b[1;32m      2\u001b[0m     set_seed(cfg\u001b[38;5;241m.\u001b[39mseed)\n\u001b[0;32m----> 3\u001b[0m     loss_fn \u001b[38;5;241m=\u001b[39m \u001b[43m_tml\u001b[49m\u001b[38;5;241m.\u001b[39mops\u001b[38;5;241m.\u001b[39mloss\u001b[38;5;241m.\u001b[39mcross_entropy_loss\n\u001b[1;32m      4\u001b[0m     reduce \u001b[38;5;241m=\u001b[39m ttml\u001b[38;5;241m.\u001b[39mops\u001b[38;5;241m.\u001b[39mReduceType\u001b[38;5;241m.\u001b[39mMEAN\n\u001b[1;32m      6\u001b[0m     causal_mask \u001b[38;5;241m=\u001b[39m build_causal_mask(cfg\u001b[38;5;241m.\u001b[39mseq_len)\n",
      "\u001b[0;31mNameError\u001b[0m: name '_tml' is not defined"
     ]
    }
   ],
   "source": [
    "def train(cfg: TrainConfig, model, optim, train_ids: np.ndarray, val_ids: np.ndarray):\n",
    "    set_seed(cfg.seed)\n",
    "    loss_fn = ttml.ops.loss.cross_entropy_loss\n",
    "    reduce = ttml.ops.ReduceType.MEAN\n",
    "\n",
    "    causal_mask = build_causal_mask(cfg.seq_len)\n",
    "    tt_mask = ttml.autograd.Tensor.from_numpy(causal_mask, ttml.Layout.TILE, ttml.autograd.DataType.BFLOAT16)   # [1,1,T,T], float32\n",
    "\n",
    "    model.train()\n",
    "    train_losses: List[float] = []\n",
    "    val_losses: List[float] = []\n",
    "\n",
    "    for step in range(1, cfg.steps + 1):\n",
    "        # ---- batch ----\n",
    "        x_u32, y_u32 = get_batch(train_ids, cfg.seq_len, cfg.batch_size)\n",
    "        # TTML shapes: inputs [B,1,1,T] (uint32), targets [B,T] (int32)\n",
    "        tt_x = ttml.autograd.Tensor.from_numpy(x_u32.reshape(cfg.batch_size, 1, 1, cfg.seq_len), ttml.Layout.ROW_MAJOR, ttml.autograd.DataType.UINT32)\n",
    "        tt_y = ttml.autograd.Tensor.from_numpy(y_u32, ttml.Layout.ROW_MAJOR, ttml.autograd.DataType.UINT32)  # [B,T] uint32\n",
    "        #print(tt_mask.to_numpy())\n",
    "        # ---- forward/backward ----\n",
    "        optim.zero_grad()\n",
    "        logits = model(tt_x, tt_mask)   \n",
    "    \n",
    "        loss = loss_fn(logits, tt_y, reduce)\n",
    "        loss.backward(False)\n",
    "        ttml.autograd.AutoContext.get_instance().reset_graph()\n",
    "        optim.step()\n",
    "\n",
    "        train_losses.append(float(loss.to_numpy()))\n",
    "\n",
    "        # ---- occasional eval on val set ----\n",
    "        \n",
    "        if (step % cfg.eval_every) == 0 or step == 1:\n",
    "            model.eval()\n",
    "            with np.errstate(all='ignore'):\n",
    "                vx_u32, vy_u32 = get_batch(val_ids, cfg.seq_len, cfg.batch_size)\n",
    "                #vtt_x = ttml.autograd.Tensor.from_numpy(vx_u32.reshape(cfg.batch_size, 1, 1, cfg.seq_len), ttml.Layout.ROW_MAJOR, ttml.autograd.DataType.UINT32)\n",
    "                #vtt_y = ttml.autograd.Tensor.from_numpy(vy_u32, ttml.Layout.ROW_MAJOR, ttml.autograd.DataType.UINT32)\n",
    "                #vlogits = model(vtt_x, tt_mask)\n",
    "    \n",
    "                #vloss = loss_fn(vlogits, vtt_y, reduce)\n",
    "                #val_losses.append(float(vloss.to_numpy()))\n",
    "                val_losses.append(train_losses[-1])  # hack to avoid eval for now\n",
    "            \n",
    "            model.train()\n",
    "\n",
    "            print(f\"step {step:5d} | train_loss {train_losses[-1]:.4f} | val_loss {val_losses[-1]:.4f}\")\n",
    "        \n",
    "\n",
    "    return train_losses, val_losses\n",
    "\n",
    "train_losses, val_losses = train(cfg, model, optim, train_ids, val_ids)\n",
    "len(train_losses), len(val_losses)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7568630",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loss curve\n",
    "plt.figure()\n",
    "plt.plot(np.arange(1, len(train_losses)+1), train_losses, label=\"train\")\n",
    "if len(val_losses) > 0:\n",
    "    # approximate x-locations for val points\n",
    "    x_val = np.linspace(1, len(train_losses), num=len(val_losses))\n",
    "    plt.plot(x_val, val_losses, marker=\"o\", linestyle=\"--\", label=\"val\")\n",
    "plt.xlabel(\"Step\")\n",
    "plt.ylabel(\"Cross-Entropy Loss\")\n",
    "plt.title(\"NanoGPT on Shakespeare (char-level)\")\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Greedy sampling (for a quick sanity check)\n",
    "def sample_greedy(model, decode_fn, start: str, max_new_tokens: int, seq_len: int):\n",
    "    # simple char-level sampling with a running window\n",
    "    model.eval()\n",
    "    # build running context from start prompt (encode to uint32)\n",
    "    ctx_ids = [ord(c) for c in start]  # placeholder; replaced next line when using our tokenizer\n",
    "    # We don't know the tokenizer here; reconstruct via decode/encode roundtrip from dataset\n",
    "    # A safe approach: pass through the training pipeline's tokenizer again.\n",
    "    # Here, we simply assume `decode` is paired with the same encode used earlier,\n",
    "    # so we'll re-encode using the train text mapping:\n",
    "    # Since we don't have encode() now, we carry the last batch as context.\n",
    "    # For a concise demo we'll just start from spaces if empty:\n",
    "    if len(start) == 0:\n",
    "        start = \" \"\n",
    "    # Using last validation chunk as seed context for correctness\n",
    "    # (In practice you'd keep the tokenizer encode() callable; omitted to keep demo minimal.)\n",
    "    text_seed = start\n",
    "\n",
    "    # Small helper: keep an internal list of token ids we feed\n",
    "    # For simplicity, use the most recent validation chunk as initial tokens\n",
    "    x_u32, _ = get_batch(val_ids, seq_len, 1)   # [1, T]\n",
    "    running = x_u32[0].tolist()\n",
    "\n",
    "    # Warm-start by replacing first few tokens with our prompt, if short\n",
    "    prompt_ids = running[:]\n",
    "    # No direct encode() handle here; we’ll just trust the random seed batch\n",
    "\n",
    "    for _ in range(max_new_tokens):\n",
    "        inp = np.array(running[-seq_len:], dtype=np.uint32).reshape(1,1,1,seq_len)\n",
    "        tt_inp = ttml.autograd.Tensor.from_numpy(inp)\n",
    "        causal_mask = build_causal_mask(seq_len)\n",
    "        tt_mask = ttml.autograd.Tensor.from_numpy(causal_mask)\n",
    "\n",
    "        logits = model(tt_inp, tt_mask)\n",
    "\n",
    "        # logits -> numpy [1, seq_len, vocab] or similar; take last position\n",
    "        np_logits = logits.to_numpy()\n",
    "        last_logits = np_logits.reshape(1, -1, vocab_size)[:, -1, :]  # [1, V]\n",
    "        next_id = int(np.argmax(last_logits, axis=-1)[0])\n",
    "        running.append(next_id)\n",
    "        ttml.autograd.AutoContext.get_instance().reset_graph()\n",
    "\n",
    "    # Decode only the newly generated tail for display (best-effort)\n",
    "    gen_ids = running[-max_new_tokens:]\n",
    "    print(\"\\nGenerated (greedy):\\n\")\n",
    "    print(textwrap.fill(decode(gen_ids), width=100))\n",
    "\n",
    "# Quick sample (greedy)\n",
    "sample_greedy(model, decode, start=\"ROMEO:\\n\", max_new_tokens=300, seq_len=cfg.seq_len)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
